package cn.edu.sjtu.stap.sootbranchinstr;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.Map.Entry;

import cn.edu.sjtu.stap.sootbranchinstr.postprocess.SwitchTemplateWithData;
import cn.edu.sjtu.stap.sootbranchinstr.template.SwitchTarget;
import cn.edu.sjtu.stap.sootbranchinstr.template.SwitchTemplate;

public class BranchRecorder {
	public static String ifRecordFileName;
	public static String switchRecordFileName;
	public static String polyRecordFileName;

	public static final String DATA_SEPARATOR = "::";
	
	private static Map<String, Map<String, String>> ifMap = new HashMap<String, Map<String, String>>();
	private static Map<String, Map<String, String>> switchMap = new HashMap<String, Map<String, String>>();
	private static Map<String, Map<String, String>> polyMap = new HashMap<String, Map<String, String>>();
	
	private static Map<String, SwitchTemplate> switchKey2Temp = new HashMap<String, SwitchTemplate>();
	
	public static void initSwitchTemplate(String tempFilePath){
		try{
			BufferedReader stReader = new BufferedReader(new FileReader(tempFilePath));
			String stLine = stReader.readLine();
			while(stLine != null){
				String[] tempSegs = stLine.split(BranchRecorder.DATA_SEPARATOR);
				SwitchTemplate st = new SwitchTemplate();
				st.setKey(tempSegs[0] + BranchRecorder.DATA_SEPARATOR + tempSegs[1]);
				List<SwitchTarget> tList = new ArrayList<SwitchTarget>();
				for(int i = 2; i < tempSegs.length; i++){
					SwitchTarget sTarget = new SwitchTarget();
					sTarget.setTargetIndex(tempSegs[i]);
					sTarget.setTargetLine(tempSegs[++i]);
					tList.add(sTarget);
				}
				st.setTargets(tList);
				
				switchKey2Temp.put(st.getKey(), st);
				stLine = stReader.readLine();
			}
			stReader.close();
		}catch(IOException ex){
			ex.printStackTrace();
		}
	}
	
	public static void dumpData(){
		try{
			if(ifRecordFileName != null){
				PrintWriter ifPw = new PrintWriter(new BufferedWriter(new FileWriter(ifRecordFileName)));
				Set<Entry<String, Map<String, String>>> ifEntrySet = ifMap.entrySet();
				for(Entry<String, Map<String, String>> ifEntry : ifEntrySet){
					StringBuilder entryStr = new StringBuilder();
					entryStr.append(ifEntry.getKey());
					Map<String, String> ifData = ifEntry.getValue();
					Set<Entry<String, String>> dataEntrySet = ifData.entrySet();
					for(Entry<String, String> dataEntry : dataEntrySet){
						entryStr.append(DATA_SEPARATOR);
						entryStr.append(dataEntry.getKey());
						entryStr.append(DATA_SEPARATOR);
						entryStr.append(dataEntry.getValue());
					}
					ifPw.println(entryStr.toString());
				}
				ifPw.close();
			}
			
			if(switchRecordFileName != null){
				PrintWriter switchPw = new PrintWriter(new BufferedWriter(new FileWriter(switchRecordFileName)));
				Set<Entry<String, Map<String, String>>> switchEntrySet = switchMap.entrySet();
				for(Entry<String, Map<String, String>> switchEntry : switchEntrySet){
					StringBuilder entryStr = new StringBuilder();
					entryStr.append(switchEntry.getKey());
					Map<String, String> switchData = switchEntry.getValue();
					Set<Entry<String, String>> dataEntrySet = switchData.entrySet();
					for(Entry<String, String> dataEntry : dataEntrySet){
						entryStr.append(DATA_SEPARATOR);
						entryStr.append(dataEntry.getKey());
						entryStr.append(DATA_SEPARATOR);
						entryStr.append(dataEntry.getValue());
					}
					switchPw.println(entryStr.toString());
				}
				switchPw.close();
			}
			
			if(polyRecordFileName != null){
				PrintWriter polyPw = new PrintWriter(new BufferedWriter(new FileWriter(polyRecordFileName)));
				Set<Entry<String, Map<String, String>>> polyEntrySet = polyMap.entrySet();
				for(Entry<String, Map<String, String>> polyEntry : polyEntrySet){
					StringBuilder entryStr = new StringBuilder();
					entryStr.append(polyEntry.getKey());
					Map<String, String>	polyData = polyEntry.getValue();
					Set<Entry<String, String>> dataEntrySet = polyData.entrySet();
					for(Entry<String, String> dataEntry : dataEntrySet){
						entryStr.append(DATA_SEPARATOR);
						entryStr.append(dataEntry.getKey());
						entryStr.append(DATA_SEPARATOR);
						entryStr.append(dataEntry.getValue());
					}
					polyPw.println(entryStr.toString());
				}
				polyPw.close();
			}
		}catch(IOException ex){
			ex.printStackTrace();
		}
	}
		
	/**
	 * Here we temporarily suppose that the combination of 
	 * method signature and line number can uniquely identify
	 * a branch statement, such as an if statement or a switch
	 * statement.
	 * 
	 * @param methodSig
	 * @param lineNum
	 */
	public static void recordIfStmt(String methodSig, String lineNum, String direction){
		String ifKey = methodSig + DATA_SEPARATOR + lineNum;
		
//		System.out.println("if: " + ifKey);
		
		Map<String, String> currentMap = ifMap.get(ifKey);
		if(currentMap == null){
			currentMap = new HashMap<String, String>();
			ifMap.put(ifKey, currentMap);
		}
		
		String directCount = currentMap.get(direction);
		if(directCount == null){
			currentMap.put(direction, "1");
		}else{
			int count = Integer.parseInt(directCount);
			currentMap.put(direction, String.valueOf(count + 1));
		}
	}
	
	public static void recordSwitchStmt(String methodSig, String lineNum, int keyValueInt){
		String keyValue = String.valueOf(keyValueInt);
		String switchKey = methodSig + DATA_SEPARATOR + lineNum;
		
		SwitchTemplate template = switchKey2Temp.get(switchKey);
		if(template == null){
			System.err.println("No template can be found for the switch statement!");
			return;
		}
		
		boolean isKeyDefault = true;
		List<SwitchTarget> tgList = template.getTargets();
		for(SwitchTarget tgt : tgList){
			if(!tgt.getTargetIndex().equals("default")){
				if(tgt.getTargetIndex().equals(keyValue)){
					isKeyDefault = false; // the current key is not in the default case.
					break;
				}
			}
		}
		
		Map<String, String> currentMap = switchMap.get(switchKey);
		if(currentMap == null){
			currentMap = new HashMap<String, String>();
			switchMap.put(switchKey, currentMap);
		}
		
		if(!isKeyDefault){
			String keyCount = currentMap.get(keyValue);
			if(keyCount == null){
				currentMap.put(keyValue, "1");
			}else{
				int count = Integer.parseInt(keyCount);
				currentMap.put(keyValue, String.valueOf(count + 1));
			}
		}else{ // when the key is in the default case
			String keyCount = currentMap.get("default");
			if(keyCount == null){
				currentMap.put("default", "1");
			}else{
				int count = Integer.parseInt(keyCount);
				currentMap.put("default", String.valueOf(count + 1));
			}
		}
	}
	
	public static void recordPolyCall(String methodSig, String lineNum, String className, String message){
		String polyKey = methodSig + DATA_SEPARATOR + lineNum + DATA_SEPARATOR + message;
		
//		System.out.println("poly: " + polyKey);
		Map<String, String> currentMap = polyMap.get(polyKey);
		if(currentMap == null){
			currentMap = new HashMap<String, String>();
			polyMap.put(polyKey, currentMap);
		}
		
		String messageKey = className;
		String polyCount = currentMap.get(messageKey);
		if(polyCount == null){
			currentMap.put(messageKey, "1");
		}else{
			int count = Integer.parseInt(polyCount);
			currentMap.put(messageKey, String.valueOf(count + 1));
		}
	}
}
